from utils.DorisQueryTool import DorisQueryTool
from TrainModel import trainModel
from Prediction import *
import os
import sys
from utils.MysqlQueryTool import  MysqlQueryTool



if __name__ == "__main__":
    # 查询所有收费站近两年数据
    print("Received arguments:", sys.argv)
    # 提取具体参数
    roadId = sys.argv[1] if len(sys.argv) > 1 else "G0004130010"
    stationId = sys.argv[2] if len(sys.argv) > 2 else "G0004130010150"
    # roadId = sys.argv[1] if len(sys.argv) > 1 else None
    # stationId = sys.argv[2] if len(sys.argv) > 2 else None
    portalIds = sys.argv[3] if len(sys.argv) > 3 else "'G000113001004020010','G000113001004020020'"
    type = sys.argv[4] if len(sys.argv) > 4 else "0"
    modelPath = sys.argv[5] if len(sys.argv) > 5 else "E:\\data\\aaa.pth"
    startTime = sys.argv[6] if len(sys.argv) > 6 else "2024-01-01 00:00:00"
    endTime = sys.argv[7] if len(sys.argv) > 7 else "2025-06-15 00:00:00"

    train_ratio = sys.argv[8] if len(sys.argv) > 8 else "70"
    val_ratio = sys.argv[9] if len(sys.argv) > 9 else "15"
    test_ratio = sys.argv[10] if len(sys.argv) > 10 else "15"
    batch_size = sys.argv[11] if len(sys.argv) > 11 else "28"
    input_length = sys.argv[12] if len(sys.argv) > 12 else "28"
    output_length = sys.argv[13] if len(sys.argv) > 13 else "7"
    learning_rate = sys.argv[14] if len(sys.argv) > 14 else "0.001"
    num_blocks = sys.argv[15] if len(sys.argv) > 15 else "2"
    dim = sys.argv[16] if len(sys.argv) > 16 else "256"
    scalerPath = sys.argv[17] if len(sys.argv) > 17 else "E:\\data\\aaa.joblib"
    epochs = sys.argv[18] if len(sys.argv) > 18 else "120"

    print("Received roadId:", roadId)
    print("Received stationId:", stationId)
    print("Received portalIds:", portalIds)
    print("Received type:", type)
    print("Received modelPath:", modelPath)
    print("Received startTime:", startTime)
    print("Received endTime:", endTime)

    print("Received train_ratio:", train_ratio)
    print("Received val_ratio:", val_ratio)
    print("Received test_ratio:", test_ratio)
    print("Received batch_size:", batch_size)
    print("Received input_length:", input_length)
    print("Received output_length:", output_length)
    print("Received learning_rate:", learning_rate)
    print("Received num_blocks:", num_blocks)
    print("Received dim:", dim)

    type = int(type)
    param = {}
    param["train_ratio"] = int(train_ratio)/100
    param["val_ratio"] = int(val_ratio)/100
    param["test_ratio"] = int(test_ratio)/100
    param["batch_size"] = int(batch_size)
    param["input_length"] = int(input_length)
    param["output_length"] = int(output_length)
    param["learning_rate"] = float(learning_rate)
    param["num_blocks"] = int(num_blocks)
    param["dim"] = int(dim)
    param["epochs"] = int(epochs)

    # abnormal_dates = ['2024-01-01', '2024-01-02', '2024-01-03', '2024-01-02', '2024-02-09', '2024-02-10', '2024-02-11',
    #                   '2024-02-12', '2024-02-13', '2024-02-14', '2024-02-15', '2024-02-16', '2024-02-17', '2024-02-18'
    #     , '2024-04-03', '2024-04-04', '2024-04-05', '2024-04-06', '2024-04-07'
    #     , '2024-04-30', '2024-05-01', '2024-05-02', '2024-05-03', '2024-05-04', '2024-05-05', '2024-05-06'
    #     , '2024-06-09', '2024-06-10', '2024-06-11', '2024-06-12'
    #     , '2024-09-14', '2024-09-15', '2024-09-16', '2024-09-17', '2024-09-18'
    #     , '2024-09-30', '2024-10-01', '2024-10-02', '2024-10-03', '2024-10-04', '2024-10-05', '2024-10-06',
    #                   '2024-10-07', '2024-10-08'
    #     , '2024-12-31', '2025-01-01', '2025-01-02'
    #     , '2025-01-27', '2025-01-28', '2025-01-29', '2025-01-30', '2025-01-31', '2025-02-01', '2025-02-02',
    #                   '2025-02-03', '2025-02-04', '2025-02-05'
    #     , '2025-04-03', '2025-04-04', '2025-04-05', '2025-04-06', '2025-04-07'
    #     , '2025-04-12', '2025-04-13'
    #     , '2025-04-30', '2025-05-01', '2025-05-02', '2025-05-03', '2025-05-04', '2025-05-05', '2025-05-06'
    #     , '2025-05-30', '2025-05-31', '2025-06-01', '2025-06-02', '2025-06-03']
    #abnormal_dates = None

    # 查询节假日
    with MysqlQueryTool() as tool:
        # 示例1：返回DataFrame
        holiday = tool.execute_query("SELECT * FROM `holiday` where is_holiday = 1 order BY date  ")
    # 查询路况
    with MysqlQueryTool() as tool:
        # 示例1：返回DataFrame
        stationInfo = tool.execute_query("SELECT status station_status,total_hour station_hour,road_id,station_id,date FROM `accident_station_info` where road_id = '"+roadId+"' and station_id = '" + stationId + "' and date >= '"+startTime+"' and date <= '"+endTime+"' order BY date  ")
        if isinstance(stationInfo, tuple):
            stationInfo = pd.DataFrame(stationInfo)
        print(stationInfo)

    with DorisQueryTool() as tool:
        if type == 0:
            # 查询入口站
            df = tool.execute_query("SELECT count( 1 ) flow,EN_TOLL_SECTION_ID roadId,EN_TOLL_STATION_ID stationId,DATE_FORMAT( EN_TIME, '%Y-%m-%d' ) date FROM `toll_en` WHERE EN_TIME >= '" + startTime + "' AND EN_TIME < '" + endTime + "' AND EN_TOLL_SECTION_ID = '" + roadId + "' AND EN_TOLL_STATION_ID = '" + stationId + "' GROUP BY roadId,stationId,date ")
            group = df.sort_values('date')
            # result = df.sort_values('date').groupby('stationId', sort=False)
            # for stationId, group in result:
            #     # 每一个收费站训练一个模型，并保存
            print("group", group)
            trainModel(group, modelPath, param,holiday,stationInfo,scalerPath)
        if type == 1:
            # 查询出口站
            if roadId is None or stationId is None:
                df = tool.execute_query(
                    "SELECT count(1) flow,DATE_FORMAT(EX_TIME, '%Y-%m-%d') date FROM `toll_ex` where EX_TIME >= '" + startTime + "' and EX_TIME < '" + endTime + "'   group by date")
                group = df.sort_values('date')
            else:
                df = tool.execute_query("SELECT count(1) flow,EX_TOLL_SECTION_ID roadId,EX_TOLL_STATION_ID stationId,EX_TOLL_STATION_NAME stationName,DATE_FORMAT(EX_TIME, '%Y-%m-%d') date FROM `toll_ex` where EX_TIME >= '"+startTime+"' and EX_TIME < '"+endTime+"' and EX_TOLL_SECTION_ID = '"+roadId+"' and EX_TOLL_STATION_ID = '"+stationId+"'  group by EX_TOLL_SECTION_ID,EX_TOLL_STATION_ID,EX_TOLL_STATION_NAME ,date")
                group = df.sort_values('date')
            # result = df.sort_values('date').groupby('stationId', sort=False)
            # for stationId, group in result:
            #     # 每一个收费站训练一个模型，并保存
            print("group", group)
            trainModel(group, modelPath,param,holiday,stationInfo,scalerPath)
        if type == 2 or type == 3:
            # 查询出省门架
            df = tool.execute_query(
                "SELECT count( 1 ) flow,TOLL_SECTION_ID roadId,DATE_FORMAT( TRANS_TIME, '%Y-%m-%d' ) date FROM	`gantry_trade` WHERE TRANS_TIME >= '" + startTime + "' AND TRANS_TIME < '" + endTime + "' AND TOLL_SECTION_ID = '" + roadId + "' AND TOLL_GANTRY_ID IN ( "+ portalIds +") GROUP BY roadId,date")
            group = df.sort_values('date')
            #     # 每一个收费单元训练一个模型，并保存
            print("group", group)

            trainModel(group, modelPath, param,holiday,stationInfo,scalerPath)